package com.lagou.gateway.standard.handler;

import com.lagou.gateway.BaseResponseBodyEntity;
import com.lagou.gateway.ResponseCodeEnum;
import com.lagou.gateway.standard.DefaultClientResponseAdapter;
import com.lagou.gateway.standard.formatter.ResponseBodyFormatter;
import com.lagou.gateway.util.ResponseBodyUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;


public class ResponseBodyHandlers {
    private static Logger LOGGER = LoggerFactory.getLogger(ResponseBodyHandlers.class);

    public static ResponseHandlerFunction<DataBuffer> getJsonBodyHandler(ResponseBodyFormatter responseBodyFormatter) {
        return (resp, body) -> {
            DefaultClientResponseAdapter clientResponseAdapter
                    = new DefaultClientResponseAdapter(body, resp.getHeaders());
            Mono<String> bodyMono = clientResponseAdapter.bodyToMono(String.class);
            return bodyMono
                    .flatMap(bodyStr -> {
                        BaseResponseBodyEntity entity
                                = responseBodyFormatter.jsonBodyFormat(resp.getStatusCode(), bodyStr);
                        return ResponseBodyUtil.writeBuffer(resp, entity);
                    })
                    .onErrorResume(e -> {
                        LOGGER.error("Response error.", e);
                        BaseResponseBodyEntity entity =
                                new BaseResponseBodyEntity(ResponseCodeEnum.SERVER_ERROR, "server error");
                        resp.setStatusCode(HttpStatus.OK);
                        return ResponseBodyUtil.writeBuffer(resp, entity);
                    })
                    .then();
        };
    }

    public static ResponseHandlerFunction<DataBuffer> getErrorHeaderBodyHandler() {
        return ((resp, body) -> {
            Flux.from(body).map(dataBuffer -> DataBufferUtils.release(dataBuffer));
            return Mono.just(resp)
                    .map(r -> {
                        BaseResponseBodyEntity entity =
                                new BaseResponseBodyEntity(ResponseCodeEnum.STATUS_UNKNOWN, "unknown status");
                        if (HttpStatus.UNAUTHORIZED.equals(resp.getStatusCode())) {
                            entity =
                                    new BaseResponseBodyEntity(ResponseCodeEnum.UNAUTHORIZED, "unauthorized");
                        } else if (resp.getStatusCode().is5xxServerError()) {
                            entity =
                                    new BaseResponseBodyEntity(ResponseCodeEnum.SERVER_ERROR, "server error");
                        } else if (HttpStatus.SEE_OTHER.equals(resp.getStatusCode())) {
                            entity =
                                new BaseResponseBodyEntity(ResponseCodeEnum.FREQUENTLY, "频繁请求，请求已被拒绝");
                        }

                        resp.setStatusCode(HttpStatus.OK);
                        return entity;
                    })
                    .flatMap(entity -> ResponseBodyUtil.writeBuffer(resp, entity))
                    .then();
        });
    }

    public static ResponseHandlerFunction<DataBuffer> getNormalBodyHandler() {
        return (resp, body) -> resp.writeWith(body);
    }
}
